package org.codetranslators.common;

import org.codetranslators.compiler.intermrep.AssemblyInstruction;
//import org.codetranslators.compiler.liveness.InterferenceGraph;
//import org.codetranslators.compiler.liveness.Node;

import java.util.Map;
import java.util.Set;
import java.util.Vector;
import java.util.HashSet;
import java.util.HashMap;
import java.util.Stack;
import java.util.Iterator;

public class RegisterAllocator 
{
	private InterferenceGraph interferenceGraph;
	private  Vector instructionList;
	private Vector outputInstructionList;
	
	private Vector colors;
	
	private int numberOfRegisters;
	
	
	public RegisterAllocator(InterferenceGraph interferenceGraph, Vector instructionList, Vector colors)
	{
		this.interferenceGraph = interferenceGraph;
		this.instructionList = instructionList;
		this.colors = colors;
		numberOfRegisters = colors.size();
		
		allocate();
	}
	
	private void allocate()
	{		
		Stack selectStack  = new Stack();
		HashMap selectedTempVsAdjTemps = new HashMap();
		HashMap coalescedVsList = new HashMap();   // A map of coalesed temp vs. it's set of corresponding children
		
		Vector igNodes = interferenceGraph.nodes();
		HashSet spilledNodes = new HashSet();
		
		// Analyze the interference graph now
		int numIter = 1;
		while(igNodes.size() != 0)
		{		
			Node igNode = (Node) igNodes.elementAt(0);
			Temp tmp = interferenceGraph.getTemp(igNode);
			System.out.println("**************ITERATION " + numIter + " ******************************");
			System.out.println("Considering temp = " + tmp);
			
			// This node has not been considered yet; process it	
			int degree = interferenceGraph.getDegree(igNode);   // Get the degree of the node
			if(degree <= numberOfRegisters -1)
			{
				processLowDegreeNode(igNodes, igNode, selectStack, selectedTempVsAdjTemps, coalescedVsList);
			}
			else
			{
				// This node is of significant degree; Get the best node from this node to the end of the list
				// of nodes
				Node bestNextNode = getBestNextNode(igNodes);
				int deg = interferenceGraph.getDegree(bestNextNode);
				if(deg <= numberOfRegisters -1)
				{
					processLowDegreeNode(igNodes, bestNextNode, selectStack, selectedTempVsAdjTemps, coalescedVsList);
				}
				else
				{
					// All nodes from igNode to the end of the list are significant degree nodes; 
					// select the best next node for spilling; if the spill priority of the next best node
					// is lesser than the spill priority of the current node being considered, make the
					// next best node the current node.
					
					int spForNextBestNode = Integer.parseInt(interferenceGraph.getSpillCost(bestNextNode));
					int spForCurrNode = Integer.parseInt(interferenceGraph.getSpillCost(igNode));
					
					if(spForCurrNode > spForNextBestNode) bestNextNode = igNode;
					
					// Remove any moves associated with this node (freeze the moves)
					removeMovesForNode(bestNextNode);
					
					// Select for spilling
					pushOnToStack(bestNextNode, selectStack,  selectedTempVsAdjTemps);
					interferenceGraph.removeNode(bestNextNode);  // This will remove all edges of this node
					igNodes.remove(bestNextNode);
					spilledNodes.add(bestNextNode);
				}
			}
			
			//System.out.println("CFG AT END OF ITERATION:" );
			//interferenceGraph.showInterferenceGraph();
			//interferenceGraph.showMoveRelatedNodes();
			
			// Reset the work-list of nodes (populate from nodes in the interference graph)
			igNodes = interferenceGraph.nodes();
			
			numIter++;
			
		}
		
		// Assign the colors now
		HashMap assignment = assignColors(selectStack, selectedTempVsAdjTemps, spilledNodes, coalescedVsList);
		
		if(assignment != null)
		{
			printAssignment(assignment);
			populateOutputInstructionList(assignment);
		}
		else
		{
			// A coloring was not possible; try to spill, rewrite code, construct cfg, ig and then analyze again.
		}
	}	

	private void pushOnToStack(Node igNode, Stack selectStack, HashMap selectedTempVsAdjTemps)
	{
		Temp temp = interferenceGraph.getTemp(igNode);
		HashSet adjList = interferenceGraph.getAdjacentNodeList(igNode);
		
		HashSet adjTemps = new HashSet();
		if(adjList != null)
		{
			Iterator iter = adjList.iterator();
			while(iter.hasNext())
			{
				Node adjNode = (Node) iter.next();
				Temp adjTemp = interferenceGraph.getTemp(adjNode);
				adjTemps.add(adjTemp);
			}
		}
		
		selectedTempVsAdjTemps.put(temp, adjTemps);
		selectStack.push(temp);
	}
	
	private void processLowDegreeNode(Vector igNodes, Node igNode, Stack selectStack,
									  HashMap selectedTempVsAdjTemps, HashMap coalescedVsList)
	{
		// This node has a degree lesser than that of the physical register of the m/c
		HashSet moveRelatedNodes = interferenceGraph.getMoveRelatedNodes(igNode);
		if(moveRelatedNodes == null || moveRelatedNodes.size() == 0)
		{
			// Not move related; remove this node from the graph and put it on the 
			// selection stack
			System.out.println("Not move related, pushing onto the stack: " + interferenceGraph.getTemp(igNode));
			pushOnToStack(igNode, selectStack,  selectedTempVsAdjTemps);
			interferenceGraph.removeNode(igNode);  // This will remove all edges of this node
			igNodes.remove(igNode);
		}
		else
		{
			// This node is move related; check if it can be coalesced with each of it's
			// move related neighbors
			Iterator iter = moveRelatedNodes.iterator();
			boolean coalesced = false;
			while(iter.hasNext())
			{
				// Coalesce with the first move related node, if possible.
				Node moveNeighbor = (Node) iter.next();
				
				// The node should not have an interference edge with this neighbor; if it does,
				// it cannot be coalesced
				HashSet igNeighbors = interferenceGraph.getAdjacentNodeList(igNode);
				if(igNeighbors != null && igNeighbors.contains(moveNeighbor)) continue;
				
				HashSet unionOfNeighbors = getUnionOfNeighbors(igNode, moveNeighbor);
				if(unionOfNeighbors.size() <= numberOfRegisters -1)
				{
					// The two nodes can be coalesced; coalesce
					coalesce(igNode, moveNeighbor, unionOfNeighbors, igNodes, coalescedVsList);
					coalesced = true;
					break;
				}
			}
			
			if(!coalesced)
			{
				// This node cannot be coalesced with any of it's move related nodes;
				// freeze it and continue. It involves disabling moves between
				// the two nodes
				iter = moveRelatedNodes.iterator();
				Vector nodesToBeRemoved = new Vector();
				while(iter.hasNext())
					nodesToBeRemoved.addElement(iter.next());
				for(int a = 0; a < nodesToBeRemoved.size(); a++)
					removeMoves(igNode, (Node) nodesToBeRemoved.elementAt(a));
				
				// Now that the moves have been removed for this node, it can be removed
				// from the interference graph
				pushOnToStack(igNode, selectStack,  selectedTempVsAdjTemps);
				interferenceGraph.removeNode(igNode);
				igNodes.remove(igNode);
			}
		}
	}
	
	/**
	 * This method returns the union of the neigbors of the input nodes. It uses HashSet objects
	 * so there are no duplicates
	 * @param igNode
	 * @param moveNeighbor
	 * @return
	 */
	private HashSet getUnionOfNeighbors(Node igNode, Node moveNeighbor)
	{		
		
		Temp t1 = interferenceGraph.getTemp(igNode);
		Temp t2 = interferenceGraph.getTemp(moveNeighbor);
		
		HashSet igNodeAdjList = interferenceGraph.getAdjacentNodeList(igNode);
		HashSet moveNeighAdjList = interferenceGraph.getAdjacentNodeList(moveNeighbor);
		
		HashSet union = new HashSet();
		
		if(igNodeAdjList == null)
		{
			// This node does not have any intereferences; check the move-related node's neighbors
			if(moveNeighAdjList == null) return union;   // The neighbor does not have any interferences
			
			Iterator iter1 = moveNeighAdjList.iterator();
			while(iter1.hasNext())
			{
				Node moveNeighNeigh = (Node) iter1.next();
				if(moveNeighNeigh != igNode) union.add(moveNeighNeigh); 
			}
		}
		else
		{
			Iterator iter = igNodeAdjList.iterator();
			while(iter.hasNext())
				union.add((Node) iter.next()); 
			
			if(moveNeighAdjList != null)
			{
				Iterator iter1 = moveNeighAdjList.iterator();
				while(iter1.hasNext())
				{
					Node moveNeighNeigh = (Node) iter1.next();
					if(moveNeighNeigh != igNode) union.add(moveNeighNeigh); 
				}
			}
		}
		return union;
	}
	
	/** This method will coalesce the two input nodes and creates a new node in the graph by 
	 * using the set that is the union of the neighbors of each of the nodes (the third parameter).
	 * 
	 * @param node
	 */
	
	private void coalesce(Node node, Node moveNeighbor, HashSet unionOfNeighbors, Vector igNodes,
			HashMap coalescedVsList)
	{
		// Create a new node in the interference graph; add this to the end of the nodes list
		// of the interference graph
		Temp temp = new Temp();
		Node newNode = interferenceGraph.newInterferenceGraphNode(temp);
		
		// Populate the map of coalesced nodes with the new node vs. it's children
		Temp[] children = new Temp[2];
		children[0] = interferenceGraph.getTemp(node);
		children[1] = interferenceGraph.getTemp(moveNeighbor);
		
		System.out.println("Coalescing " + children[0] + " and " + children[1] + " into " + temp);
		
		coalescedVsList.put(temp, children);
		
		// Add edges from the new node to each element in the union set
		Iterator iter = unionOfNeighbors.iterator();
		while(iter.hasNext())
		{
			Node neighbor = (Node) iter.next();
			interferenceGraph.updateAdjacencyMatrix(newNode, neighbor);
		}
		
		// Remove the move between the nodes, since they are being coalesed
		removeMoves(node, moveNeighbor);
		
		// Create moves for the new node if any; the moves will be the union of the moves related
		// to the two nodes
		HashSet union = new HashSet();
		HashSet movesForNode = interferenceGraph.getMoveRelatedNodes(node);
		HashSet movesForMoveNeigh =  interferenceGraph.getMoveRelatedNodes(moveNeighbor);
		Iterator iter1 = movesForNode.iterator();
		while(iter1.hasNext())
		{
			Node nd = (Node) iter1.next();
			Temp t1 = interferenceGraph.getTemp(nd);
			union.add(nd);
			//if(t1 != null) union.add(nd);
		}
		
		Iterator iter2 = movesForMoveNeigh.iterator();
		while(iter2.hasNext())
		{
			Node nd = (Node) iter2.next();
			Temp t2 = interferenceGraph.getTemp(nd);
			union.add(nd);
			//if(t2 != null) union.add(nd);
		}
		
		Iterator iter3 = union.iterator();
		while(iter3.hasNext())
		{
			interferenceGraph.addToMoveRelatedNodes(newNode, (Node)iter3.next());
		}
		
		// The spill cost of the new coalesced node will be the maximum of the spill cost of the two nodes
		// that are being coalesced
		int spCostOfNewNode = Integer.parseInt(interferenceGraph.getSpillCost(node));
		int spCostOfMoveNeigh = Integer.parseInt(interferenceGraph.getSpillCost(moveNeighbor));
		if(spCostOfNewNode < spCostOfMoveNeigh ) spCostOfNewNode = spCostOfMoveNeigh;
		interferenceGraph.setSpillCost(newNode, ""+spCostOfNewNode);
		
		// Remove the two nodes from the interference graph
		interferenceGraph.removeNode(node);
		interferenceGraph.removeNode(moveNeighbor); 
		igNodes.remove(node);
		igNodes.remove(moveNeighbor);
		removeMovesForNode(node);
		removeMovesForNode(moveNeighbor);
	}
	
	private void removeMoves(Node node, Node moveNeighbor)
	{
		HashSet moveRelatedNodesForInputNode = interferenceGraph.getMoveRelatedNodes(node);
		HashSet moveRelatedNodesForMoveNeigh = interferenceGraph.getMoveRelatedNodes(moveNeighbor);
		moveRelatedNodesForInputNode.remove(moveNeighbor);
		moveRelatedNodesForMoveNeigh.remove(node);
	}
	
	/**
	 * This method will return the next "best" node to be processed while iterating through the nodes of the interference
	 * graph. The next "best" node will be either a low-degree node or a high degree node with a high spill cost.
	 * 
	 * @return
	 */
	
	private Node getBestNextNode(Vector igNodes)
	{
		//	Get the next node with insignificant degree; during the sequentional search,
		// find the nodes of high degree with the highest spill cost - this will
		// help if there are no nodes with low degree
		Node nextNodeWithLowerDegree = null;
		Node highDegNodeWithMaxSpillCost = null;
		int maxSpillCost = 0;
		int numNodes = igNodes.size();
		for(int a = 1; a < numNodes; a++)
		{
			Node nextIgNode = (Node) igNodes.elementAt(a);  // Get the next node
			Temp tmp = interferenceGraph.getTemp(nextIgNode);
			//int degreeOfNextNode = nextIgNode.degree();   // Get the degree of this node
			int degreeOfNextNode = interferenceGraph.getDegree(nextIgNode);   // Get the degree of this node
			if(degreeOfNextNode <= numberOfRegisters -1) 
			{
				nextNodeWithLowerDegree = nextIgNode;
				break;
			}
			else
			{
				int spCost = Integer.parseInt(interferenceGraph.getSpillCost(nextIgNode));
				if(spCost > maxSpillCost) 
				{
					maxSpillCost = spCost;
					highDegNodeWithMaxSpillCost = nextIgNode;
				}
			}
		}
		
		// Check if a lower degree node was found; if found, return it; else return the high-degree node 
		// with maximum spill cost
		if(nextNodeWithLowerDegree != null) return nextNodeWithLowerDegree;
		
		// No low degree node was found; return the one that was found with max spill cost
		return highDegNodeWithMaxSpillCost;
	}
	
	private void removeMovesForNode(Node node)
	{
		HashSet moveRelatedNodes = interferenceGraph.getMoveRelatedNodes(node);
		if(moveRelatedNodes != null && moveRelatedNodes.size() != 0)
		{
			Vector nodesToBeRemoved = new Vector();
			Iterator iter = moveRelatedNodes.iterator();
			while(iter.hasNext())
				nodesToBeRemoved.addElement(iter.next());
			
			for(int a = 0; a < nodesToBeRemoved.size(); a++)
				removeMoves(node, (Node)nodesToBeRemoved.elementAt(a));
		}
	}
	
	private HashMap assignColors(Stack selectStack, HashMap selectedTempVsAdjTemps, HashSet spilledNodes,
								 HashMap coalescedVsList)
	{
		interferenceGraph.reset();
		
		HashSet usedColors = new HashSet();
		
		HashMap tempVsNewNode = new HashMap();
		HashMap tempVsColor = new HashMap(); 
		
		while(!selectStack.isEmpty())
		{
			Temp temp = (Temp) selectStack.pop();
			
			Node newNode = interferenceGraph.newInterferenceGraphNode(temp);
			tempVsNewNode.put(temp, newNode);
			
			HashSet adjTemps = (HashSet) selectedTempVsAdjTemps.get(temp);
			
			// Find out the colors that CANNOT be used for this new node; these will be the 
			// colors of the adjacent nodes of the new temp
			Iterator iter = adjTemps.iterator();
			HashSet prohibitedColorsForNewNode = new HashSet();
			while(iter.hasNext())
			{
				Temp adjTemp = (Temp) iter.next();
				prohibitedColorsForNewNode.add(tempVsColor.get(adjTemp));
			}
			
			if(prohibitedColorsForNewNode.size() >= numberOfRegisters)
			{
				// No color available for new node; have to spill (a potential spill has now
				// become an actual spill) and re-try
				return null;
			}
			
			// Must try to use an already used color; iterate through the usedColors set and find 
			// if there is one such color
			Iterator iter1 = usedColors.iterator();
			String colorToUse = null;
			while(iter1.hasNext())
			{
				String usedColor = (String) iter1.next();
				if(!prohibitedColorsForNewNode.contains(usedColor))
				{
					colorToUse = usedColor;
					break;
				}
			}
			
			if(colorToUse == null)
			{
				colorToUse = (String) colors.remove(0); // Get the first non-used color
			}
				
			tempVsColor.put(temp, colorToUse);
			
			// If temp is a coalesced temporary, add it's children also to the tempVsColor map
			HashSet coalescedTemps = getCoalescedTemps(temp, coalescedVsList);
			Iterator coalTempIter = coalescedTemps.iterator();
			while(coalTempIter.hasNext())
				tempVsColor.put((Temp) coalTempIter.next(), colorToUse);
			
			usedColors.add(colorToUse);
		}
		return tempVsColor;
	}
	
	private HashSet getCoalescedTemps(Temp temp, HashMap coalescedVsList)
	{
		HashSet coalescedTemps = new HashSet();
		
		Temp[] children = (Temp[]) coalescedVsList.get(temp);
		if(children == null) return coalescedTemps;
		
		coalescedTemps.add(children[0]);
		coalescedTemps.add(children[1]);
		coalescedTemps.addAll(getCoalescedTemps(children[0], coalescedVsList));
		coalescedTemps.addAll(getCoalescedTemps(children[1], coalescedVsList));
		
		return coalescedTemps;
	}
	
	private void populateOutputInstructionList(HashMap assignment)
	{
		outputInstructionList = new Vector();
		
		// Replace temp variable with the actual physical register names in the code
		int numIns = instructionList.size();
		for(int i = 0; i < numIns; i++)
		{
			AssemblyInstruction ins = (AssemblyInstruction) instructionList.elementAt(i);
			String insStr = ins.getAssemCode(false);
			StringBuffer insStrBuff = new StringBuffer(insStr);
			
			Vector def = ins.getDestList();
			Vector use = ins.getSrcList();
			
			if(ins.isMove())
			{
				// Get the color assigned to each temp (for the dest and the use temp)
				String colorDest = (String) assignment.get((Temp) def.elementAt(0));  
				String colorUse = (String) assignment.get((Temp) use.elementAt(0));
				// If they are the same, ignore them; this will help get rid of unnecessary move instructions
				if(colorDest == colorUse) continue;
				
			}
			
			if(def != null) replaceStr(def, insStr, assignment,insStrBuff );
			insStr = insStrBuff.toString();
			if(use != null) replaceStr(use, insStr, assignment,insStrBuff );
			
			outputInstructionList.addElement(insStrBuff.toString());				
		}
	}
	
	private void replaceStr(Vector vec, String insStr, HashMap assignment, StringBuffer insStrBuff )
	{
		for(int j = 0; j < vec.size();j++)
		{
			Temp tmp = (Temp) vec.elementAt(j);
			int startIndex = insStr.indexOf(tmp.toString());
			int len = tmp.toString().length();
			String color = (String) assignment.get(tmp);
			if(color != null) insStrBuff.replace(startIndex , startIndex + len, color);
			insStr = insStrBuff.toString();
		}
	}
	
	private void printDegrees()
	{
		System.out.println("In printDegrees");
		String str = "";
		for(int i = 0; i < interferenceGraph.nodes().size(); i++)
		{
			Node node = (Node) interferenceGraph.nodes().elementAt(i);
			if(node == null) continue;
			Temp temp = (Temp) interferenceGraph.getTemp(node);
			if(temp == null) continue;
			str = temp.toString() + " degree : " + interferenceGraph.getDegree(node);
			System.out.println(str);
		}
	}
	
	private void printAssignment(HashMap assignment)
	{
		HashMap colorVsTemps = new HashMap();
		
		System.out.println("COLOR ASSIGNMENT: ");
		Set entries = assignment.entrySet();
		Iterator iter = entries.iterator();
		while(iter.hasNext())
		{
			Map.Entry entry = (Map.Entry) iter.next();
			Temp temp = (Temp)entry.getKey();
			String color = (String) entry.getValue();
			
			HashSet tempList = (HashSet) colorVsTemps.get(color);
			if(tempList == null) tempList = new HashSet();
			tempList.add(temp);
			colorVsTemps.put(color, tempList);
		}
		
		entries = colorVsTemps.entrySet();
		iter = entries.iterator();
		while(iter.hasNext())
		{
			Map.Entry entry = (Map.Entry) iter.next();
			String color = (String) entry.getKey();
			
			System.out.print(color + ": ");
			
			HashSet tempList = (HashSet) entry.getValue();
			Iterator iter1 = tempList.iterator();
			String tmpListStr = "";
			while(iter1.hasNext())
			{
				Temp tmp = (Temp) iter1.next();
				if(tmp != null)
					tmpListStr += tmp.toString() + " ";
				else
					tmpListStr += "???" + " ";
			}
			System.out.println(tmpListStr);
		}
	}

	public Vector getOutputInstructionList()
	{
		return outputInstructionList;
	}
	
	public int getNumRegistersUsed()
	{
		return numberOfRegisters - colors.size();
	}
}
